package com.mdframe.ai.agent.tools;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;

/**
 * @date 2024/10/17
 */
@Service
@Slf4j
public class MySQLSchemaService {
    
    private static String URL;
    
    private static String USERNAME;
    
    private static String PASSWORD;
    
    static {
        URL = "jdbc:mysql://127.0.0.1:3306/forge_admin";
        USERNAME = "root";
        PASSWORD = "yaomindong";
    }
    
    public Map<String, String> getTableSchema(String tableName) {
        try {
            String driver = "com.mysql.jdbc.Driver";
            Class.forName(driver);
            try (Connection connection = DriverManager.getConnection(URL, USERNAME, PASSWORD);
                    Statement statement = connection.createStatement()) {
                String sql = "DESCRIBE " + tableName;
                ResultSet resultSet = statement.executeQuery(sql);
                Map<String, String> schemaMap = new HashMap<>();
                while (resultSet.next()) {
                    String columnName = resultSet.getString("Field");
                    String columnType = resultSet.getString("Type");
                    schemaMap.put(columnName, columnType);
                }
                return schemaMap;
            }
        } catch (Exception e) {
            log.error("获取表结构失败", e);
            return null;
        }
    }
    
}
